package io.kiki.sba.registry.util;


import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class AtomicMap<K, V> {

    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock();
    private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock();
    private Map<K, V> data = new ConcurrentHashMap<>();

    public V put(K key, V value) {
        readLock.lock();
        try {
            return data.put(key, value);
        } finally {
            readLock.unlock();
        }
    }

    public Map<K, V> getAndReset() {
        writeLock.lock();
        try {
            Map<K, V> ret = data;
            data = new ConcurrentHashMap<>();
            return ret;
        } finally {
            writeLock.unlock();
        }
    }
}
